from modelscope.metainfo import Trainers
from modelscope.trainers import build_trainer

kwargs = dict(
    model='./cv_tinynas_object-detection_damoyolo',
    gpu_ids=[0, 1],  # 使用的GPU设备ID列表
    batch_size=16,  # 根据显存调整（AdamW对batch_size更鲁棒）
    max_epochs=500,  # 总训练轮次

    # 优化器配置（AdamW参数）
    optimizer='AdamW',  # 使用AdamW优化器
    base_lr=0.001,  # AdamW的初始学习率通常比SGD小（建议0.001-0.0001）
    weight_decay=0.05,  # AdamW通常需要更高的weight_decay（推荐0.05）
    betas=(0.9, 0.999),  # AdamW的动量参数（默认值）
    amsgrad=False,  # 是否使用AMSGrad变体

    # 禁用数据增强
    # mosaic_prob=0.0,  # 关闭Mosaic增强
    # mixup_prob=0.0,  # 关闭Mixup增强
    # hsv_h=0.0,  # 禁用色调变换
    # hsv_s=0.0,  # 禁用饱和度变换
    # hsv_v=0.0,  # 禁用明度变换
    # flip_prob=0.0,  # 禁用水平翻转

    # 学习率调度
    lr_schedule='CosineAnnealing',  # 余弦退火调度
    warmup_iters=500,  # 学习率预热步数

    # 模型保存与验证
    ckpt_interval_epochs=40,  # 每5个epoch保存一次模型
    save_best_checkpoint=True,  # 保存最佳模型
    eval_interval_epochs=2,  # 每1个epoch验证一次

    # 数据集配置
    num_classes=2,
    train_image_dir='./datasets/coco_format_dataset/images',
    val_image_dir='./datasets/coco_format_dataset/images',
    train_ann='./datasets/coco_format_dataset/annotations/instance_train.json',
    val_ann='./datasets/coco_format_dataset/annotations/instance_val.json',

    # 其他配置
    work_dir='./workdirs',
    log_interval=50,
    pretrained=True,  # 使用预训练权重
    resume_from=None,  # 可指定checkpoint路径继续训练
)

trainer = build_trainer(
    name=Trainers.tinynas_damoyolo,
    default_args=kwargs
)
if __name__ == '__main__':
    trainer.train()


